"""可视化调试模块
包含矩阵热力图、特征谱可视化等功能
"""
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Tuple, List, Optional, Union, Dict
from matplotlib.patches import Circle
from mpl_toolkits.mplot3d import Axes3D
import warnings

class MatrixVisualization:
    @staticmethod
    def plot_matrix_heatmap(matrix: np.ndarray, 
                          title: str = 'Matrix Heatmap',
                          cmap: str = 'viridis',
                          annotate: bool = False,
                          figsize: Tuple[int, int] = (8, 6)) -> None:
        """
        绘制矩阵热力图
        :param matrix: 输入矩阵
        :param title: 图表标题
        :param cmap: 颜色映射
        :param annotate: 是否显示数值标注
        :param figsize: 图形大小
        """
        plt.figure(figsize=figsize)
        
        if annotate and matrix.size <= 100:  # 只在小矩阵上显示数值
            sns.heatmap(matrix, annot=True, cmap=cmap, fmt='.2f')
        else:
            plt.imshow(matrix, cmap=cmap, aspect='auto')
            plt.colorbar()
        
        plt.title(title)
        plt.xlabel('Column Index')
        plt.ylabel('Row Index')
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def plot_sparse_matrix(matrix: np.ndarray, 
                          title: str = 'Sparse Matrix Structure',
                          figsize: Tuple[int, int] = (8, 6)) -> None:
        """
        绘制稀疏矩阵结构图
        :param matrix: 输入矩阵
        :param title: 图表标题
        :param figsize: 图形大小
        """
        plt.figure(figsize=figsize)
        
        # 找到非零元素
        rows, cols = np.nonzero(matrix)
        
        plt.scatter(cols, rows, s=1, c='black', marker='s')
        plt.xlim(-0.5, matrix.shape[1] - 0.5)
        plt.ylim(matrix.shape[0] - 0.5, -0.5)  # 反转y轴
        plt.xlabel('Column Index')
        plt.ylabel('Row Index')
        plt.title(f'{title} (Sparsity: {1 - len(rows)/matrix.size:.2%})')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()

    @staticmethod
    def plot_eigenvalues(eigenvalues: np.ndarray, 
                        title: str = 'Eigenvalue Spectrum',
                        figsize: Tuple[int, int] = (8, 6)) -> None:
        """
        绘制特征值谱
        :param eigenvalues: 特征值数组
        :param title: 图表标题
        :param figsize: 图形大小
        """
        plt.figure(figsize=figsize)
        
        real_parts = np.real(eigenvalues)
        imag_parts = np.imag(eigenvalues)
        
        # 绘制特征值点
        plt.scatter(real_parts, imag_parts, color='red', marker='o', s=50, alpha=0.7)
        
        # 添加单位圆（用于稳定性分析）
        circle = Circle((0, 0), 1, fill=False, color='blue', linestyle='--', alpha=0.5)
        plt.gca().add_patch(circle)
        
        # 添加坐标轴
        plt.axhline(0, color='black', linestyle='-', linewidth=0.5, alpha=0.7)
        plt.axvline(0, color='black', linestyle='-', linewidth=0.5, alpha=0.7)
        
        plt.xlabel('Real Part')
        plt.ylabel('Imaginary Part')
        plt.title(title)
        plt.grid(True, alpha=0.3)
        plt.axis('equal')
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def plot_singular_values(singular_values: np.ndarray,
                           title: str = 'Singular Value Spectrum',
                           log_scale: bool = True,
                           figsize: Tuple[int, int] = (8, 6)) -> None:
        """
        绘制奇异值谱
        :param singular_values: 奇异值数组
        :param title: 图表标题
        :param log_scale: 是否使用对数刻度
        :param figsize: 图形大小
        """
        plt.figure(figsize=figsize)
        
        indices = np.arange(1, len(singular_values) + 1)
        
        if log_scale:
            plt.semilogy(indices, singular_values, 'bo-', markersize=4)
            plt.ylabel('Singular Values (log scale)')
        else:
            plt.plot(indices, singular_values, 'bo-', markersize=4)
            plt.ylabel('Singular Values')
        
        plt.xlabel('Index')
        plt.title(title)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def plot_matrix_decomposition(original: np.ndarray, 
                                 factors: List[np.ndarray],
                                 decomp_type: str = 'SVD',
                                 figsize: Tuple[int, int] = (15, 5)) -> None:
        """
        可视化矩阵分解结果
        :param original: 原始矩阵
        :param factors: 分解因子列表
        :param decomp_type: 分解类型
        :param figsize: 图形大小
        """
        n_factors = len(factors)
        fig, axes = plt.subplots(1, n_factors + 1, figsize=figsize)
        
        # 绘制原始矩阵
        im0 = axes[0].imshow(original, cmap='viridis', aspect='auto')
        axes[0].set_title('Original Matrix')
        axes[0].set_xlabel('Column')
        axes[0].set_ylabel('Row')
        plt.colorbar(im0, ax=axes[0])
        
        # 绘制分解因子
        factor_names = {
            'SVD': ['U', 'Σ', 'V^T'],
            'LU': ['L', 'U'],
            'QR': ['Q', 'R'],
            'Cholesky': ['L']
        }
        
        names = factor_names.get(decomp_type, [f'Factor {i+1}' for i in range(n_factors)])
        
        for i, (factor, name) in enumerate(zip(factors, names)):
            if factor.ndim == 1:  # 对角矩阵的情况
                factor_matrix = np.diag(factor)
            else:
                factor_matrix = factor
            
            im = axes[i+1].imshow(factor_matrix, cmap='viridis', aspect='auto')
            axes[i+1].set_title(f'{name}')
            axes[i+1].set_xlabel('Column')
            axes[i+1].set_ylabel('Row')
            plt.colorbar(im, ax=axes[i+1])
        
        plt.suptitle(f'{decomp_type} Decomposition')
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def plot_condition_number_analysis(matrices: List[np.ndarray],
                                     labels: List[str],
                                     title: str = 'Condition Number Analysis',
                                     figsize: Tuple[int, int] = (10, 6)) -> None:
        """
        分析并可视化矩阵条件数
        :param matrices: 矩阵列表
        :param labels: 标签列表
        :param title: 图表标题
        :param figsize: 图形大小
        """
        condition_numbers = []
        
        for matrix in matrices:
            try:
                cond_num = np.linalg.cond(matrix)
                condition_numbers.append(cond_num)
            except:
                condition_numbers.append(np.inf)
        
        plt.figure(figsize=figsize)
        
        # 使用对数刻度
        plt.semilogy(range(len(condition_numbers)), condition_numbers, 'bo-', markersize=8)
        
        # 添加标签
        plt.xticks(range(len(labels)), labels, rotation=45)
        plt.ylabel('Condition Number (log scale)')
        plt.title(title)
        plt.grid(True, alpha=0.3)
        
        # 添加警告线
        plt.axhline(y=1e12, color='red', linestyle='--', alpha=0.7, label='Ill-conditioned threshold')
        plt.legend()
        
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def plot_convergence_history(errors: List[float],
                               title: str = 'Convergence History',
                               log_scale: bool = True,
                               figsize: Tuple[int, int] = (8, 6)) -> None:
        """
        绘制算法收敛历史
        :param errors: 误差历史
        :param title: 图表标题
        :param log_scale: 是否使用对数刻度
        :param figsize: 图形大小
        """
        plt.figure(figsize=figsize)
        
        iterations = range(1, len(errors) + 1)
        
        if log_scale:
            plt.semilogy(iterations, errors, 'b-', linewidth=2, marker='o', markersize=4)
            plt.ylabel('Error (log scale)')
        else:
            plt.plot(iterations, errors, 'b-', linewidth=2, marker='o', markersize=4)
            plt.ylabel('Error')
        
        plt.xlabel('Iteration')
        plt.title(title)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def plot_tensor_slices(tensor: np.ndarray,
                          slice_dim: int = 0,
                          max_slices: int = 9,
                          figsize: Tuple[int, int] = (12, 8)) -> None:
        """
        可视化张量切片
        :param tensor: 输入张量
        :param slice_dim: 切片维度
        :param max_slices: 最大显示切片数
        :param figsize: 图形大小
        """
        if tensor.ndim < 3:
            raise ValueError("Tensor must have at least 3 dimensions")
        
        n_slices = min(tensor.shape[slice_dim], max_slices)
        cols = int(np.ceil(np.sqrt(n_slices)))
        rows = int(np.ceil(n_slices / cols))
        
        fig, axes = plt.subplots(rows, cols, figsize=figsize)
        if rows == 1 and cols == 1:
            axes = [axes]
        elif rows == 1 or cols == 1:
            axes = axes.flatten()
        else:
            axes = axes.flatten()
        
        for i in range(n_slices):
            if slice_dim == 0:
                slice_data = tensor[i, :, :]
            elif slice_dim == 1:
                slice_data = tensor[:, i, :]
            else:
                slice_data = tensor[:, :, i]
            
            im = axes[i].imshow(slice_data, cmap='viridis', aspect='auto')
            axes[i].set_title(f'Slice {i}')
            axes[i].set_xticks([])
            axes[i].set_yticks([])
            plt.colorbar(im, ax=axes[i])
        
        # 隐藏多余的子图
        for i in range(n_slices, len(axes)):
            axes[i].set_visible(False)
        
        plt.suptitle(f'Tensor Slices (dimension {slice_dim})')
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def plot_numerical_stability(matrix: np.ndarray,
                               perturbation_levels: List[float],
                               operation: str = 'eigenvalues',
                               figsize: Tuple[int, int] = (10, 6)) -> None:
        """
        分析数值稳定性
        :param matrix: 输入矩阵
        :param perturbation_levels: 扰动水平列表
        :param operation: 操作类型
        :param figsize: 图形大小
        """
        plt.figure(figsize=figsize)
        
        # 计算原始结果
        if operation == 'eigenvalues':
            original_result = np.linalg.eigvals(matrix)
            original_metric = np.max(np.real(original_result))
        elif operation == 'condition':
            original_metric = np.linalg.cond(matrix)
        else:
            raise ValueError(f"Unsupported operation: {operation}")
        
        relative_errors = []
        
        for eps in perturbation_levels:
            # 添加随机扰动
            perturbation = eps * np.random.randn(*matrix.shape)
            perturbed_matrix = matrix + perturbation
            
            try:
                if operation == 'eigenvalues':
                    perturbed_result = np.linalg.eigvals(perturbed_matrix)
                    perturbed_metric = np.max(np.real(perturbed_result))
                elif operation == 'condition':
                    perturbed_metric = np.linalg.cond(perturbed_matrix)
                
                relative_error = abs(perturbed_metric - original_metric) / abs(original_metric)
                relative_errors.append(relative_error)
            except:
                relative_errors.append(np.inf)
        
        plt.loglog(perturbation_levels, relative_errors, 'bo-', markersize=6)
        plt.xlabel('Perturbation Level')
        plt.ylabel('Relative Error')
        plt.title(f'Numerical Stability Analysis - {operation.title()}')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def create_matrix_report(matrix: np.ndarray, 
                           matrix_name: str = 'Matrix',
                           save_path: Optional[str] = None) -> Dict[str, Union[float, str]]:
        """
        生成矩阵分析报告
        :param matrix: 输入矩阵
        :param matrix_name: 矩阵名称
        :param save_path: 保存路径
        :return: 分析报告字典
        """
        report = {
            'name': matrix_name,
            'shape': matrix.shape,
            'dtype': str(matrix.dtype),
            'size': matrix.size,
            'memory_usage_mb': matrix.nbytes / (1024 * 1024)
        }
        
        # 基本统计
        report['min_value'] = float(np.min(matrix))
        report['max_value'] = float(np.max(matrix))
        report['mean_value'] = float(np.mean(matrix))
        report['std_value'] = float(np.std(matrix))
        
        # 稀疏性
        zero_elements = np.count_nonzero(matrix == 0)
        report['sparsity'] = zero_elements / matrix.size
        
        # 方阵特性
        if matrix.shape[0] == matrix.shape[1]:
            report['is_square'] = True
            
            # 对称性
            report['is_symmetric'] = np.allclose(matrix, matrix.T)
            
            # 条件数
            try:
                report['condition_number'] = float(np.linalg.cond(matrix))
            except:
                report['condition_number'] = np.inf
            
            # 行列式
            try:
                report['determinant'] = float(np.linalg.det(matrix))
            except:
                report['determinant'] = np.nan
            
            # 迹
            report['trace'] = float(np.trace(matrix))
            
            # 秩
            try:
                report['rank'] = int(np.linalg.matrix_rank(matrix))
            except:
                report['rank'] = -1
        else:
            report['is_square'] = False
        
        # 保存报告
        if save_path:
            with open(save_path, 'w') as f:
                for key, value in report.items():
                    f.write(f"{key}: {value}\n")
        
        return report